package site.https;

import cn.hutool.core.util.IdUtil;
import io.netty.bootstrap.Bootstrap;
import io.netty.bootstrap.ServerBootstrap;
import io.netty.buffer.PooledByteBufAllocator;
import io.netty.channel.*;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.nio.NioServerSocketChannel;
import io.netty.channel.socket.nio.NioSocketChannel;
import io.netty.handler.codec.http.*;
import io.netty.handler.timeout.IdleStateEvent;
import io.netty.handler.timeout.IdleStateHandler;
import io.netty.util.ReferenceCountUtil;
import lombok.Data;
import lombok.experimental.Accessors;
import lombok.extern.slf4j.Slf4j;

import java.net.InetSocketAddress;
import java.net.MalformedURLException;
import java.net.URL;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.TimeUnit;

/**
 * @author haocongshun
 * @date 2023/02/22 15:44:22
 */
public class HttpsProxy {

    private static final EventLoopGroup group = new NioEventLoopGroup();

    private static Map<String, Tunnel> tunnelMap = new HashMap<>();

    public static void main(String[] args) throws Exception {
        ServerBootstrap serverBootstrap = new ServerBootstrap();
        EventLoopGroup bossGroup = new NioEventLoopGroup();
        EventLoopGroup workGroup = new NioEventLoopGroup();

        serverBootstrap.group(bossGroup, workGroup)
                .channel(NioServerSocketChannel.class)
                .childOption(ChannelOption.ALLOCATOR, PooledByteBufAllocator.DEFAULT)
                .childOption(ChannelOption.TCP_NODELAY, true)
                .childHandler(new ChannelInitializer<Channel>() {
                    @Override
                    protected void initChannel(Channel ch) throws Exception {
                        ch.pipeline()
                                .addLast("idle", new IdleStateHandler(0, 0, 100, TimeUnit.SECONDS))
                                .addLast("http-decoder", new HttpRequestDecoder())
                                .addLast("http-encoder", new HttpResponseEncoder())
                                .addLast("connect-judge", new TunnelHandler())
                                .addLast("exception", new ExceptionDuplexHandler());
                    }
                });

        serverBootstrap.bind(8090).sync().channel().closeFuture().sync();
    }

    @Slf4j
    public static class ExceptionDuplexHandler extends ChannelDuplexHandler {
        @Override
        public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
            ctx.channel().close();
            log.error("{} happen error, will be close : {}", ctx.channel().id(), cause.getMessage());
            cause.printStackTrace();
        }

        @Override
        public void userEventTriggered(ChannelHandlerContext ctx, Object evt) throws Exception {
            if (evt instanceof IdleStateEvent) {
                ctx.channel().close();
                log.warn("{} idle timeout, will be close", ctx.channel().id());
            }
        }
    }

    @Slf4j
    public static class TunnelHandler extends ChannelInboundHandlerAdapter {

        @Override
        public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
            if (msg instanceof HttpRequest) {
                HttpRequest httpMsg = (HttpRequest) msg;

                if (httpMsg.method() == HttpMethod.CONNECT) {
                    InetSocketAddress dstAddress = getDstAddress(httpMsg);
                    Tunnel tunnel = new Tunnel().setHandleRequest(ctx.channel());
                    tunnelMap.put(tunnel.getId(), tunnel);
                    ChannelFuture channelFuture = connectRemote(dstAddress, new DefaultRemoteFlowTransfer(tunnel.getId(), "request"));
                    channelFuture.addListener(new ChannelFutureListener() {
                        @Override
                        public void operationComplete(ChannelFuture future) throws Exception {
                            if(!future.isSuccess()){
                                log.error("[{}-] connect [{}:{}] fail {}.", tunnel.getId(), dstAddress.getHostName(), dstAddress.getPort(), future.cause().getMessage());
                                tunnel.close();
                            }
//                            tunnel.getHandleRequest().pipeline().remove(TunnelHandler.class);
                            String pip = null;
                            try {
                                ChannelPipeline pipeline = tunnel.getHandleRequest().pipeline();
                                StringBuilder sb = new StringBuilder();
                                pipeline.forEach(handler -> sb.append(handler.getClass().getSimpleName()).append(" -> "));
                                sb.append("END");
                                pip = sb.toString();
                                tunnel.getHandleRequest().pipeline().addBefore("exception", "default-remote-flow-transfer", new DefaultRemoteFlowTransfer(tunnel.getId(), "response"));
                            }catch (Exception e){
                                log.error("pip:{}",pip,e);
                            }
                            tunnel.getHandleRequest().writeAndFlush(new DefaultFullHttpResponse(HttpVersion.HTTP_1_1, HttpResponseStatus.OK))
                                    .addListener(new ChannelFutureListener() {
                                        public void operationComplete(ChannelFuture future) {
                                            tunnel.getHandleRequest().pipeline().remove(HttpRequestDecoder.class);
                                            tunnel.getHandleRequest().pipeline().remove(HttpResponseEncoder.class);
                                        }
                                    });
                        }
                    });
                }
                ReferenceCountUtil.release(msg);
                ctx.pipeline().remove(this);
            } else {
                ctx.fireChannelRead(msg);
            }
        }

        /**
         * get address form httpMsg
         *
         * @param httpMsg httpMsg
         * @return address
         */
        private InetSocketAddress getDstAddress(HttpRequest httpMsg) {
            InetSocketAddress result = null;

            String uri = httpMsg.uri();
            if (uri.startsWith("http://") || uri.startsWith("https://")) {
                try {
                    URL url = new URL(uri);
                    result = new InetSocketAddress(url.getHost(), url.getPort() == -1 ? 80 : url.getPort());
                } catch (MalformedURLException e) {
                    throw new IllegalArgumentException(httpMsg.uri() + " is getDstAddress fail");
                }
            } else {
                String host = uri.contains(":") ? uri.substring(0, uri.lastIndexOf(":")) : uri;
                int port = uri.contains(":") ? Integer.valueOf(uri.substring(uri.lastIndexOf(":") + 1)) : 80;
                return new InetSocketAddress(host, port);
            }

            return result;
        }


        private ChannelFuture connectRemote(InetSocketAddress dstAddress, final ChannelHandler... channelHandlers) {

            Bootstrap remoteBootstrap = new Bootstrap();

            return remoteBootstrap.group(group)
                    .option(ChannelOption.SO_KEEPALIVE, true)
                    .option(ChannelOption.TCP_NODELAY, true)
                    .option(ChannelOption.CONNECT_TIMEOUT_MILLIS, 60000)
                    .channel(NioSocketChannel.class)
                    .handler(new ChannelInitializer<Channel>() {
                        protected void initChannel(Channel ch) {
                            for (ChannelHandler channelHandler : channelHandlers) {
                                ch.pipeline().addLast(channelHandler);
                            }
                        }
                    }).connect(dstAddress);
        }

    }


    @Slf4j
    public static class DefaultRemoteFlowTransfer extends ChannelInboundHandlerAdapter {

        private String tunnelId;
        private String to;

        public DefaultRemoteFlowTransfer(String id, String to) {
            this.tunnelId = id;
            this.to = to;
        }

        @Override
        public void channelActive(ChannelHandlerContext ctx) throws Exception {
            if ("request".equals(to)) {
                tunnelMap.get(tunnelId).setHandleResponse(ctx.channel());
            }
        }

        @Override
        public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
            if ("request".equals(to)) {
                tunnelMap.get(tunnelId).handleRequest.writeAndFlush(msg);
            } else {
                tunnelMap.get(tunnelId).handleResponse.writeAndFlush(msg);
            }

        }

        @Override
        public void channelInactive(ChannelHandlerContext ctx) throws Exception {
            log.info("隧道{}关闭", tunnelId);
            tunnelMap.get(tunnelId).close();
        }
    }

    @Data
    @Accessors(chain = true)
    public static class Tunnel {

        private String id;
        private Channel handleRequest;
        private Channel handleResponse;

        public Tunnel() {
            this.id = IdUtil.randomUUID();
        }

        public void close() {
            this.handleResponse.close();
            this.handleRequest.close();
        }
    }
}
